/*******************************************************************************
 * Copyright 2010 DemandSpot
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 ******************************************************************************/

package com.demandspot.aws;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.MultiThreadedHttpConnectionManager;
import org.apache.commons.httpclient.methods.GetMethod;
import org.apache.log4j.Logger;

import com.xerox.amazonws.ec2.EC2Exception;
import com.xerox.amazonws.ec2.Jec2;
import com.xerox.amazonws.ec2.LaunchConfiguration;
import com.xerox.amazonws.ec2.ReservationDescription;
import com.xerox.amazonws.ec2.TerminatingInstanceDescription;
import com.xerox.amazonws.ec2.ReservationDescription.Instance;
import com.xerox.amazonws.monitoring.MetricStatisticsResult;
import com.xerox.amazonws.monitoring.Monitoring;
import com.xerox.amazonws.monitoring.MonitoringException;
import com.xerox.amazonws.monitoring.StandardUnit;
import com.xerox.amazonws.monitoring.Statistics;

public class AWSEC2Service {
	
	private static final Logger logger = Logger.getLogger(AWSEC2Service.class);
	
	private static final long CHECK_PERIOD_MINUTES = 1;
	
	private static MultiThreadedHttpConnectionManager connManager = new MultiThreadedHttpConnectionManager();
	
	private Jec2 ec2Service;
	private ScheduledThreadPoolExecutor executer = new ScheduledThreadPoolExecutor(5);

	private String accessKey;
	private String secretKey;
		
	
	public AWSEC2Service() throws FileNotFoundException, IOException, AWSS3ServiceException {		
		this(AWSEC2Service.class.getResourceAsStream("/conf/aws.properties"));
	}	

	public AWSEC2Service(InputStream propInputStream) throws IOException {
		Properties props = new Properties();
		props.load(propInputStream);
		init(props);
	}
	
	public AWSEC2Service(File propsFile) throws IOException {
		Properties props = new Properties();
		props.load(new FileInputStream(propsFile));
		init(props);
	}
	
	public AWSEC2Service(URL propsURL) throws IOException {
		this(new File(propsURL.getFile()));
	}
	
	public AWSEC2Service(String accessKey, String secretKey) {
		Properties props = new Properties();
		props.put(AWSConsts.AWS_ACCESS_KEY, accessKey);
		props.put(AWSConsts.AWS_SECRET_KEY, secretKey);
		init(props);
	}
	
	
	private void init(Properties props) {
		accessKey = props.getProperty(AWSConsts.AWS_ACCESS_KEY);
		secretKey = props.getProperty(AWSConsts.AWS_SECRET_KEY);
		
		accessKey = System.getProperty(AWSConsts.AWS_ACCESS_KEY, accessKey);
		secretKey = System.getProperty(AWSConsts.AWS_SECRET_KEY, secretKey);
		String secureStr = props.getProperty("aws.ec2.secure");
		
		boolean secure = false;
		if (secureStr != null && !secureStr.trim().isEmpty()) {
			secure = Boolean.parseBoolean(secureStr);
		}
		
		ec2Service = new Jec2(accessKey, secretKey, secure);
	}
	
	public List<TerminatingInstanceDescription> terminateLocalInstance() throws EC2Exception {
		String instanceId = getLocalInstanceId();
		if (instanceId != null) {
			return ec2Service.terminateInstances(Collections.singletonList(instanceId));
		}
		else {
			logger.info("Not running on EC2 - sending System.exit(1) instead of terminating instance");
			System.exit(1);
			return Collections.emptyList();	
		}		
	}
	
	public List<Instance> startInstances(LaunchConfiguration launchConfiguration) throws EC2Exception {
		ReservationDescription resDesc = ec2Service.runInstances(launchConfiguration);
		List<Instance> instances = resDesc.getInstances();		
		return instances;
	}
	
	public ReservationDescription startInstancesForReservation(LaunchConfiguration launchConfiguration) throws EC2Exception {
		ReservationDescription resDesc = ec2Service.runInstances(launchConfiguration);
		return resDesc;
	}

	public Set<Instance> getInstancesRunning(List<Instance> instances) throws EC2Exception {		
		String[] ids = getInstaceIds(instances);		
		Map<String, Instance> instancesMap = new HashMap<String, Instance>();
		for (Instance instance : instances) {
			instancesMap.put(instance.getInstanceId(), instance);
		}
		
		Set<Instance> result = new HashSet<Instance>();
		List<ReservationDescription> resrvs = ec2Service.describeInstances(ids);
		for (ReservationDescription resrv : resrvs) {
			List<Instance> resrvInstances = resrv.getInstances();
			for (Instance instance : resrvInstances) {
				if (instancesMap.containsKey(instance.getInstanceId())) {
					Instance existing = instancesMap.get(instance.getInstanceId());
					if ("running".equals(instance.getState())) {
						result.add(existing);	
					}
				}
			}
		}
		
		return result;
	}
	
	public List<Instance> getRunningInstances() throws EC2Exception {
		List<Instance> res = new ArrayList<Instance>();
		List<String> params = new ArrayList<String>();
        List<ReservationDescription> instances = ec2Service.describeInstances(params);
        for (ReservationDescription reserv : instances) {
	        if (reserv.getInstances() != null) {
                for (Instance inst : reserv.getInstances()) {
                	res.add(inst);                	
                }
	        }
        }
        return res;
	}

	public Set<TerminatingInstanceDescription> getInstancesTerminated(List<TerminatingInstanceDescription> instances) throws EC2Exception {		
		String[] ids = getTerminsatedInstaceIds(instances);		
		Map<String, TerminatingInstanceDescription> instancesMap = new HashMap<String, TerminatingInstanceDescription>();
		for (TerminatingInstanceDescription instance : instances) {
			instancesMap.put(instance.getInstanceId(), instance);
		}
		
		Set<TerminatingInstanceDescription> result = new HashSet<TerminatingInstanceDescription>();
		List<ReservationDescription> resrvs = ec2Service.describeInstances(ids);
		for (ReservationDescription resrv : resrvs) {
			List<Instance> resrvInstances = resrv.getInstances();
			for (Instance instance : resrvInstances) {
				if (instancesMap.containsKey(instance.getInstanceId())) {
					TerminatingInstanceDescription existing = instancesMap.get(instance.getInstanceId());
					if ("terminated".equals(instance.getState())) {
						result.add(existing);	
					}
				}
			}
		}
		
		return result;
	}
	
	public List<Instance> startInstancesBlocking(LaunchConfiguration launchConfiguration) throws EC2Exception {
		logger.info("Waiting for EC2 instance to start");		
		final Semaphore sem = new Semaphore(-1);
		
		ReservationDescription reservationDescription = startInstancesForReservation(launchConfiguration);
		final List<Instance> instances =  reservationDescription.getInstances();
		
		List<Instance> toReturn = new ArrayList<Instance>(instances);
		executer.scheduleAtFixedRate(new Runnable() {
			public void run() {
				try {
					Set<Instance> updated = getInstancesRunning(instances);
					instances.removeAll(updated);
				} catch (EC2Exception e) {
					logger.error(e);
				}
				
				if (instances.isEmpty()) {
					executer.remove(this);
					sem.release();					
				}
			}			
		}, CHECK_PERIOD_MINUTES, CHECK_PERIOD_MINUTES, TimeUnit.MINUTES);

		try {
			sem.acquire();
		} catch (InterruptedException e) {
			logger.error(e);
		}
		
		String reservationId = reservationDescription.getReservationId();
		List<ReservationDescription> descs = ec2Service.describeInstances(new ArrayList<String>());
		for (ReservationDescription reservationDescription2 : descs) {
			if (reservationDescription2.getReservationId().equals(reservationId)) {
				toReturn = reservationDescription2.getInstances();
			}
		}
		
		logger.info("EC2 instance started");
		return toReturn;
	}
	
	
	public List<TerminatingInstanceDescription> terminateInstances(List<Instance> instances, final ShudownCallback shudownCallback) throws EC2Exception {
		String[] ids = getInstaceIds(instances);
		final List<TerminatingInstanceDescription> terminating =  ec2Service.terminateInstances(ids);
		executer.scheduleAtFixedRate(new Runnable() {
			public void run() {
				try {
					Set<TerminatingInstanceDescription> updated = getInstancesTerminated(terminating);					
					terminating.removeAll(updated);
				} catch (EC2Exception e) {
					logger.error(e);					
				}
				
				if (terminating.isEmpty()) {
					executer.remove(this);
					shudownCallback.onShutdown();
				}
			}			
		}, CHECK_PERIOD_MINUTES, CHECK_PERIOD_MINUTES, TimeUnit.MINUTES);		
		return terminating;
	}

	public void terminateInstance(String... instanceIds) throws EC2Exception {
		ec2Service.terminateInstances(instanceIds); 
	}
	
	public Instance describeInstance(String instanceId) throws EC2Exception {
		List<ReservationDescription> resrvs = ec2Service.describeInstances(new String[]{instanceId});
		return resrvs.get(0).getInstances().get(0);
	}
	
	private String[] getInstaceIds(List<Instance> instances) {
		String[] ids = new String[instances.size()];
		for (int i = 0; i < ids.length; i++) {
			ids[i] = instances.get(i).getInstanceId();
		}
		return ids;
	}

	private String[] getTerminsatedInstaceIds(List<TerminatingInstanceDescription> instances) {
		String[] ids = new String[instances.size()];
		for (int i = 0; i < ids.length; i++) {
			ids[i] = instances.get(i).getInstanceId();
		}
		return ids;
	}
	
	public Map<Instance, MetricStatisticsResult> getCPUUtilization(Date fromDate, Date toDate) throws MonitoringException, EC2Exception {
		Map<Instance, MetricStatisticsResult> res = new HashMap<Instance, MetricStatisticsResult>();
				
		Monitoring m = new Monitoring(accessKey, secretKey);		
		List<Statistics> stats = new ArrayList<Statistics>();
		stats.add(Statistics.AVERAGE);
				
		List<Instance> running = getRunningInstances();
		for (Instance instance : running) {
			Map<String, String> dimensions = new HashMap<String, String>();
			dimensions.put("InstanceId", instance.getInstanceId());			
			MetricStatisticsResult metricStatisticsResult = m.getMetricStatistics(60, stats, "AWS/EC2", dimensions, fromDate, toDate, "CPUUtilization", StandardUnit.PERCENT, null);
			if (!metricStatisticsResult.getDatapoints().isEmpty()) {
				res.put(instance, metricStatisticsResult);
			}
		}
		
		return res;
	}
	
	public static String getLocalInstanceId() {
		return getLocalMachineMetaData("instance-id");		
	}
	
	private static String getLocalMachineMetaData(String metaDataKey) {
        HttpClient client = new HttpClient(connManager);
        GetMethod getMethod = new GetMethod("http://169.254.169.254/latest/meta-data/" + metaDataKey);
        
        try {
            client.executeMethod(getMethod);
        } catch (IOException ex) {        	
		} finally {
			getMethod.releaseConnection();
        }
		
		try {
			return getMethod.getResponseBodyAsString();
		} catch (IOException e) {
			return "";
		}
	}
}
